// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/partition_streaming_window_build.h>

namespace kumo::pollux::exec {
    PartitionStreamingWindowBuild::PartitionStreamingWindowBuild(
        const std::shared_ptr<const core::WindowNode> &windowNode,
        pollux::memory::MemoryPool *pool,
        const common::SpillConfig *spillConfig,
        tsan_atomic<bool> *nonReclaimableSection)
        : WindowBuild(windowNode, pool, spillConfig, nonReclaimableSection) {
    }

    void PartitionStreamingWindowBuild::buildNextPartition() {
        partitionStartRows_.push_back(sortedRows_.size());
        sortedRows_.insert(sortedRows_.end(), inputRows_.begin(), inputRows_.end());
        inputRows_.clear();
    }

    void PartitionStreamingWindowBuild::addInput(RowVectorPtr input) {
        for (auto i = 0; i < inputChannels_.size(); ++i) {
            decodedInputVectors_[i].decode(*input->childAt(inputChannels_[i]));
        }

        for (auto row = 0; row < input->size(); ++row) {
            char *newRow = data_->newRow();

            for (auto col = 0; col < input->childrenSize(); ++col) {
                data_->store(decodedInputVectors_[col], row, newRow, col);
            }

            if (previousRow_ != nullptr &&
                compareRowsWithKeys(previousRow_, newRow, partitionKeyInfo_)) {
                buildNextPartition();
            }

            inputRows_.push_back(newRow);
            previousRow_ = newRow;
        }
    }

    void PartitionStreamingWindowBuild::noMoreInput() {
        buildNextPartition();

        // Help for last partition related calculations.
        partitionStartRows_.push_back(sortedRows_.size());
    }

    std::shared_ptr<WindowPartition>
    PartitionStreamingWindowBuild::nextPartition() {
        POLLUX_CHECK_GT(
            partitionStartRows_.size(), 0, "No window partitions available");

        ++currentPartition_;
        POLLUX_CHECK_LE(
            currentPartition_,
            partitionStartRows_.size() - 2,
            "All window partitions consumed");

        // Erase previous partition.
        if (currentPartition_ > 0) {
            const auto numPreviousPartitionRows =
                    partitionStartRows_[currentPartition_];
            data_->eraseRows(
                melon::Range<char **>(sortedRows_.data(), numPreviousPartitionRows));
            sortedRows_.erase(
                sortedRows_.begin(), sortedRows_.begin() + numPreviousPartitionRows);
            sortedRows_.shrink_to_fit();
            for (int i = currentPartition_; i < partitionStartRows_.size(); ++i) {
                partitionStartRows_[i] =
                        partitionStartRows_[i] - numPreviousPartitionRows;
            }
        }

        const auto partitionSize = partitionStartRows_[currentPartition_ + 1] -
                                   partitionStartRows_[currentPartition_];
        const auto partition = melon::Range(
            sortedRows_.data() + partitionStartRows_[currentPartition_],
            partitionSize);

        return std::make_shared<WindowPartition>(
            data_.get(), partition, inversedInputChannels_, sortKeyInfo_);
    }

    bool PartitionStreamingWindowBuild::hasNextPartition() {
        return partitionStartRows_.size() > 0 &&
               currentPartition_ <
               static_cast<vector_size_t>(partitionStartRows_.size() - 2);
    }
} // namespace kumo::pollux::exec
